"""FPS Module"""
import mindspore as ms
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore.common.tensor import Tensor
import os
import sys
import numpy as np
import pointnet2_util

# def fps_cuda():
#     return ops.Custom("sampling_gpu.so:furthest_point_sampling", out_shape=[2, 256], out_dtype=ms.int32, func_type="aot")
    
class FPSModule(nn.Cell):
    def __init__(self, num_proposal):
        super().__init__()
        self.num_proposal = num_proposal

    def construct(self, xyz, features):
        """
        Args:
            xyz: (B,K,3)
            features: (B,C,K)
        """
        # Farthest point sampling (FPS)
        transpose = ops.Transpose()
        fps_idx = Tensor(np.tile(np.arange(self.num_proposal),xyz.shape[0]).reshape(xyz.shape[0],-1), ms.int32)
        new_xyz = pointnet2_util.index_points(xyz, fps_idx)
        new_features = transpose(pointnet2_util.index_points(transpose(features, (0, 2, 1)), fps_idx), (0, 2, 1))
        return new_xyz, new_features, fps_idx
